"""
BEL 推理路径生成器

使用 vLLM 生成多条推理路径，严格按照设计文档要求
复用现有的 vLLM 推理框架和统一的工具函数
"""

import json
import os
import sys
from typing import Dict, List, Optional, Any
from pathlib import Path
import time

# 复用现有的推理模块
from vllm import LLM, SamplingParams
from inference.vllm_infer import extract_answer_try_all_methods, template, prompt, prepare_data
from PIL import Image
import jsonlines
from loguru import logger

from .utils import (
    create_llm_instance, create_sampling_params, prepare_data_custom,
    normalize_image_path, ensure_dir, count_lines, ProgressBar,
    load_config, validate_config, setup_logging
)


class RolloutGenerator:
    """
    推理路径生成器
    
    使用vLLM生成多条推理路径，为错误定位提供数据
    严格按照设计文档要求，复用统一的工具函数
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化生成器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.llm = None
        
        # 从配置中提取参数
        self.num_rollouts = config['algorithm']['num_rollouts']
        self.batch_size = config['algorithm']['batch_size']
        
        logger.info(f"RolloutGenerator 初始化完成")
    
    def _ensure_llm(self):
        """延迟初始化 LLM 实例，复用统一的初始化函数"""
        if self.llm is None:
            self.llm = create_llm_instance(self.config)
    
    def generate_rollouts(self, input_file: str, output_file: str, resume: bool = True) -> None:
        """
        生成推理路径
        
        Args:
            input_file: 输入数据文件路径
            output_file: 输出文件路径
            resume: 是否支持断点续跑
        """
        logger.info(f"开始生成推理路径: {input_file} -> {output_file}")
        
        # 确保输出目录存在
        ensure_dir(output_file)
        
        # 检查断点续跑
        start_idx = 0
        if resume and os.path.exists(output_file):
            start_idx = count_lines(output_file)
            logger.info(f"检测到已有输出文件，从第 {start_idx + 1} 条记录开始")
        
        # 加载数据
        input_data, origin_data = self._load_data(input_file)
        
        if start_idx >= len(input_data):
            logger.info("所有数据已处理完成")
            return
        
        # 过滤需要处理的数据
        input_data = input_data[start_idx:]
        origin_data = origin_data[start_idx:]
        
        logger.info(f"需要处理 {len(input_data)} 条数据")
        
        # 延迟初始化LLM
        self._ensure_llm()
        
        # 创建采样参数
        sampling_params = create_sampling_params(self.config)
        
        # 批量处理
        with ProgressBar(len(input_data), "生成推理路径") as pbar:
            with open(output_file, 'a', encoding='utf-8') as f:
                for i in range(0, len(input_data), self.batch_size):
                    batch_input = input_data[i:i + self.batch_size]
                    batch_origin = origin_data[i:i + self.batch_size]
                    
                    try:
                        # 生成推理路径
                        outputs = self.llm.generate(batch_input, sampling_params)
                        
                        # 处理输出
                        for j, (output, origin) in enumerate(zip(outputs, batch_origin)):
                            rollouts = []
                            for generated_output in output.outputs:
                                rollouts.append({
                                    'output': generated_output.text.strip()
                                })
                            
                            # 构建结果记录
                            result = {
                                **origin,  # 保留原始数据
                                'rollouts': rollouts
                            }
                            
                            # 写入文件
                            f.write(json.dumps(result, ensure_ascii=False) + '\n')
                            f.flush()
                            
                            pbar.update(1)
                            
                    except Exception as e:
                        logger.error(f"批次 {i//self.batch_size + 1} 处理失败: {e}")
                        # 跳过失败的批次，避免中断整个流程
                        for _ in range(len(batch_input)):
                            pbar.update(1)
                        continue
        
        logger.info(f"推理路径生成完成: {output_file}")
    
    def run(self):
        """
        运行生成器
        """
        data_config = self.config['data']
        output_config = self.config['output']
        
        input_file = data_config['data_path']
        output_file = output_config['rollouts']
        
        self.generate_rollouts(input_file, output_file)
    
    def _load_data(self, input_file: str):
        """
        加载输入数据
        
        Args:
            input_file: 输入文件路径
            
        Returns:
            Tuple: (input_data, origin_data)
        """
        data_config = self.config['data']
        dataset = data_config.get('dataset', 'custom')
        image_root = data_config.get('image_root', '')
        max_samples = self.config['algorithm'].get('num_samples_per_run', -1)
        
        logger.info(f"加载数据: {input_file}")
        logger.info(f"数据集类型: {dataset}")
        logger.info(f"图片根目录: {image_root}")
        
        if max_samples > 0:
            logger.info(f"限制处理样本数: {max_samples}")
        
        # 优先使用自定义数据准备函数
        try:
            input_data, origin_data = prepare_data_custom(
                dataset, input_file, image_root, max_samples
            )
            logger.info(f"使用自定义数据准备函数成功加载 {len(input_data)} 条数据")
            return input_data, origin_data
        except Exception as e:
            logger.warning(f"自定义数据准备函数失败: {e}")
            
            # 回退到标准数据准备函数
            try:
                input_data, origin_data = prepare_data(
                    dataset, input_file, image_root
                )
                # 如果需要限制样本数，在这里处理
                if max_samples > 0 and len(input_data) > max_samples:
                    input_data = input_data[:max_samples]
                    origin_data = origin_data[:max_samples]
                logger.info(f"使用标准数据准备函数成功加载 {len(input_data)} 条数据")
                return input_data, origin_data
            except Exception as e2:
                logger.error(f"标准数据准备函数也失败: {e2}")
                raise


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="BEL 推理路径生成器")
    parser.add_argument("--config", type=str, required=True, help="配置文件路径")
    parser.add_argument("--log_level", type=str, default="INFO", help="日志级别")
    
    args = parser.parse_args()
    
    # 设置日志
    setup_logging(args.log_level)
    
    # 加载配置
    config = load_config(args.config)
    if not validate_config(config):
        logger.error("配置文件验证失败")
        sys.exit(1)
    
    # 创建生成器并运行
    generator = RolloutGenerator(config)
    
    try:
        start_time = time.time()
        generator.run()
        elapsed_time = time.time() - start_time
        logger.info(f"生成完成，耗时: {elapsed_time:.2f} 秒")
    except KeyboardInterrupt:
        logger.info("用户中断，程序退出")
    except Exception as e:
        logger.error(f"生成失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main() 